#include "SolvePnP.h"

float PI = 3.1415926535f;

void SolvePnP::run()
{
    /**
     * @brief 初始化甲板四角在世界坐标下的位置
     * @brief 顺序：左上，左下，右上，右下
     * 
     */
    obj_points.push_back(Point3f(-armor_length / 2, -armor_width / 2, 0)); // tl  -> Top Left
    obj_points.push_back(Point3f(armor_length / 2, -armor_width / 2, 0));  // bl
    obj_points.push_back(Point3f(-armor_length / 2, armor_width / 2, 0));  // tr
    obj_points.push_back(Point3f(armor_length / 2, armor_width / 2, 0));   // br

    VideoCapture cap(video_path);
    /**
     * @brief 记录每帧姿态以及位置信息
     * @details 输出 到 "../record.csv" 中
     * 
     */
    ofstream outFile;
    outFile.open("../record.csv", ios::out);
    outFile.clear();
    outFile << "时间： "
            << "欧拉角                   "
            << "位置向量                 "
            <<"距离" << endl;
    auto start = std::chrono::steady_clock::now();
    float norm;
    while (cap.read(src))
    {
        contours_mask = Mat::zeros(src.size(), CV_8UC1);
        vector<Edge> edges;

        /**
         * @brief 提取 R 通道
         * 
         */
        split(src, channels);
        r_channel = channels[2];
        threshold(r_channel, threshold_mask, 100, 255, THRESH_BINARY);
        findContours(threshold_mask, contours, hierarchy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);

        // 筛选光条 同任务一
        for (auto &contour : contours)
        {
            RotatedRect min_rect = minAreaRect(contour);

            // 规整化 RotatedRect 规整结果，以竖直线上半段为极轴，顺时针为正，逆时针为负，[-90,90）
            float width = min_rect.size.width;
            float height = min_rect.size.height;
            float angle = min_rect.angle;

            /**
             * @brief 以高宽比以及面积大小筛选光条
             * 
             */
            RotatedRect reg_rect = angle < 45 ? min_rect : RotatedRect(min_rect.center, Size2f(height, width), -90 + angle);
            if ((reg_rect.size.height / reg_rect.size.width > 3) && (reg_rect.size.height / reg_rect.size.width < 8) && min_rect.size.area() > 150 && min_rect.size.area() < 3000)
            {
                Point point_edge_up = getPointbyAngleAndDistance(reg_rect.center, reg_rect.angle, reg_rect.size.height * 0.5f);
                Point point_edge_down = getPointbyAngleAndDistance(reg_rect.center, reg_rect.angle, -reg_rect.size.height * 0.5f);
                Edge edge;
                edge.up_point = point_edge_up;
                edge.down_point = point_edge_down;
                edge.center = reg_rect.center;
                edges.push_back(edge);
            }
        }
        /**
         * @brief 按照x坐标排序灯条
         * 
         */
        sort(edges.begin(), edges.end(),
             [&](const Edge &a, const Edge &b)
             {
                 return a.center.x < b.center.x;
             });

        /**
         * @brief 如果有两条边及以上时，认为甲板可能存在
         * 
         */
        if (edges.size() > 1)
        {
            for (int i = 0; i < edges.size() - 1; i++)
            {
                if (judgeArmor(edges[i], edges[i + 1], 1.2, 0.35))
                {
                    /**
                     * @brief 按同样顺序初始化摄像头中位置
                     * 
                     */
                    img_points.clear();
                    img_points.push_back(edges[i].up_point);
                    img_points.push_back(edges[i].down_point);
                    img_points.push_back(edges[i + 1].up_point);
                    img_points.push_back(edges[i + 1].down_point);

                    /**
                     * @brief 核心函数，求出旋转以及位置向量
                     * 
                     */
                    solvePnP(obj_points, img_points, camera_matrix, distortion_coefficients, rvec, tvec);
                    norm = getNorm(tvec);
                    /**
                     * @brief 求出欧拉角
                     * 
                     */
                    Rodrigues(rvec, r_mat);
                    Rodrigues(tvec, t_mat);
                    theta_x = atan2f(r_mat.at<float>(2, 1), r_mat.at<float>(2, 2));
                    theta_y = atan2f(-r_mat.at<float>(2, 0),
                                     sqrtf(r_mat.at<float>(2, 1) * r_mat.at<float>(2, 1) + r_mat.at<float>(2, 2) * r_mat.at<float>(2, 2)));
                    theta_z = atan2f(r_mat.at<float>(1, 0), r_mat.at<float>(0, 0));

                    // 转换为角度
                    theta_x = theta_x * (180 / PI);
                    theta_y = theta_y * (180 / PI);
                    theta_z = theta_z * (180 / PI);

                    cout << "欧拉角："
                         << "[" << theta_x << "," << theta_y << "," << theta_z << "]" << endl;
                    cout << "tvec: "
                         << "[" << tvec[0] << "," << tvec[1] << "," << tvec[2] << "]" << endl;
                    cout << "距离："
                         << norm<<endl;
                    // 相机世界位置
                    // Mat P = -r_mat.inv() * Mat(tvec);
                    // cout << "Position" << P <<endl;
                }
            }
        }

        namedWindow("Result", WINDOW_AUTOSIZE);
        imshow("Result", src);
        waitKey(delay);
        auto end = chrono::steady_clock::now();
        chrono::duration<double> elapsed_seconds = end - start;
        outFile << elapsed_seconds.count() << " "
                << "[" << theta_x << "," << theta_y << "," << theta_z << "]"
                << " "
                << "[" << tvec[0] << "," << tvec[1] << "," << tvec[2] << "]" 
                <<" "
                << norm
                << endl;
    }
    outFile.close();
}

/**
 * @brief 同任务一中 判断是否为甲板
 * 
 * @param a 
 * @param b 
 * @param upper_limit 
 * @param lower_limit 
 * @return true 
 * @return false 
 */
bool SolvePnP::judgeArmor(const Edge &a, const Edge &b, float upper_limit, float lower_limit)
{
    float height_1 = getDistance4points(a.up_point, a.down_point);
    float height_2 = getDistance4points(b.up_point, b.down_point);
    float height = (height_1 + height_2) / 2;
    float width = getDistance4points(a.center, b.center);
    float ratio = height / width;
    return ratio < upper_limit && ratio > lower_limit;
}

/**
 * @brief 构建对象，导入参数
 * 
 * @param path 
 */
SolvePnP::SolvePnP(const String &path)
{
    FileStorage fread(path, FileStorage::READ);
    fread["video_path"] >> video_path;
    fread["camera_matrix"] >> camera_matrix;
    fread["distortion_coefficients"] >> distortion_coefficients;
    fread["armor_width"] >> armor_width;
    fread["armor_length"] >> armor_length;
    fread["delay"] >> delay;
};